import numpy
import numpy as np
import odl
from odl.contrib import torch as odl_torch
import cv2
import astra


def get_operator():
    para_ini = MyInitialization()
    fp = build_gemotry(para_ini)
    op_mod_fp = odl_torch.OperatorModule(fp)
    op_mod_pT = odl_torch.OperatorModule(fp.adjoint)
    return op_mod_fp, op_mod_pT


## 640geo
class initialization:
    def __init__(self):
        self.param = {}
        # self.reso = 512 / 416 * 0.03
        self.reso = 512 / 364 * 0.03

        # image
        # self.param['nx_h'] = 416
        # self.param['ny_h'] = 416
        self.param['nx_h'] = 364  # modified by Gao
        self.param['ny_h'] = 364
        self.param['sx'] = self.param['nx_h'] * self.reso
        self.param['sy'] = self.param['ny_h'] * self.reso

        ## view
        self.param['startangle'] = 0
        self.param['endangle'] = 2 * np.pi

        # self.param['nProj'] = 640
        self.param['nProj'] = 720

        ## detector
        self.param['su'] = 2 * np.sqrt(self.param['sx'] ** 2 + self.param['sy'] ** 2)
        self.param['nu_h'] = 641
        # self.param['dde'] = 1075 * self.reso
        # self.param['dso'] = 1075 * self.reso
        self.param['dde'] = 1  # diameter of the detector?
        self.param['dso'] = 1  # diameter of the source?

        self.param['u_water'] = 0.192


class MyInitialization:
    def __init__(self):
        self.param = {}
        # self.reso = 512 / 416 * 0.03
        self.reso = 256 / 256 * 0.03

        # image
        # self.param['nx_h'] = 416
        # self.param['ny_h'] = 416
        self.param['nx_h'] = 256  # modified by Gao
        self.param['ny_h'] = 256
        self.param['sx'] = self.param['nx_h'] * self.reso
        self.param['sy'] = self.param['ny_h'] * self.reso

        ## view
        self.param['startangle'] = 0
        self.param['endangle'] = 2 * np.pi

        # self.param['nProj'] = 640
        self.param['nProj'] = 320

        ## detector
        self.param['su'] = 2 * np.sqrt(self.param['sx'] ** 2 + self.param['sy'] ** 2)
        self.param['nu_h'] = 197
        self.param['dde'] = 1075 * self.reso
        self.param['dso'] = 1075 * self.reso
        # self.param['dde'] = 1  # diameter of the detector?
        # self.param['dso'] = 1  # diameter of the source?

        self.param['u_water'] = 0.192


def build_gemotry(param):
    reco_space_h = odl.uniform_discr(
        min_pt=[-param.param['sx'] / 2.0, -param.param['sy'] / 2.0],
        max_pt=[param.param['sx'] / 2.0, param.param['sy'] / 2.0],
        shape=[param.param['nx_h'], param.param['ny_h']],
        dtype='float32')

    angle_partition = odl.uniform_partition(param.param['startangle'], param.param['endangle'],
                                            param.param['nProj'])

    detector_partition_h = odl.uniform_partition(-(param.param['su'] / 2.0), (param.param['su'] / 2.0),
                                                 param.param['nu_h'])

    geometry_h = odl.tomo.FanBeamGeometry(angle_partition, detector_partition_h,
                                          src_radius=param.param['dso'],
                                          det_radius=param.param['dde'])

    ray_trafo_hh = odl.tomo.RayTransform(reco_space_h, geometry_h, impl='astra_cuda')
    return ray_trafo_hh


def build_gemotry_own(param):
    """
    目前锥束加入到网络中都比较难，需要的显存过多。
    因此测试使用deeplesion数据，使用扇形束即可，但是参数和InDuDoNet使用的可能不太一样。
    """
    # reco_space_h = odl.uniform_discr(
    #     min_pt=[-param.param['sx'] / 2.0, -param.param['sy'] / 2.0],
    #     max_pt=[param.param['sx'] / 2.0, param.param['sy'] / 2.0], shape=[param.param['nx_h'], param.param['ny_h']],
    #     dtype='float32')
    # Todo 还没搞懂这个参数是干啥用的，先用官方代码中的示例
    reco_space = odl.uniform_discr(
        min_pt=[-20, -20], max_pt=[20, 20], shape=[256, 256], dtype='float32')
    # startangle 0 endangle 2Pi nProj 720  # 320来自于ADN中制作deeplesion数据时使用的参数（光源旋转次数）
    angle_partition = odl.uniform_partition(0, 2 * np.pi, 320)
    # su 38.0140605566 nu_h 641  0.171来自于ADN中制作deeplesion数据时matlab函数para2fan源码中gammaRad变量
    detector_partition_h = odl.uniform_partition(-0.171, 0.171, 197)
    # dso 39.69230769230769  dde 39.69230769230769
    geometry_h = odl.tomo.FanBeamGeometry(angle_partition,
                                          detector_partition_h,
                                          src_radius=1075,  # Radius of the source circle
                                          det_radius=1075)  # Radius of the detector circle

    ray_trafo_hh = odl.tomo.RayTransform(reco_space, geometry_h, impl='astra_cuda')
    return ray_trafo_hh


class AstraReconstruction():
    def __init__(self, config):
        self.vol_geom = astra.create_vol_geom(config.size, config.size)
        # self.vol_geom = astra.create_vol_geom(256, 256)
        # self.proj_geom = astra.create_proj_geom('parallel', 1.0, 196, np.linspace(0, 2 * np.pi, 320, False))
        # proj_geom = astra_create_proj_geom('fanflat', det_width, det_count, angles, source_origin, origin_det);
        self.proj_geom = astra.create_proj_geom('parallel', 1.0, 384, np.linspace(0, 2 * np.pi, 360, False))

        # self.proj_geom = astra.create_proj_geom('fanflat', 1.0, 196, np.linspace(0, 2 * np.pi, 360, False), 1000, 1000)

    def do_fbp(self, sinogram):
        # debug_result
        rec_id = astra.data2d.create('-vol', self.vol_geom)
        sinogram_id = astra.data2d.create('-sino', self.proj_geom, sinogram)
        # create configuration
        cfg = astra.astra_dict('FBP_CUDA')
        cfg['ReconstructionDataId'] = rec_id
        cfg['ProjectionDataId'] = sinogram_id
        cfg['option'] = {'FilterType': 'shepp-logan'}

        # possible values for FilterType:
        # none, ram-lak, shepp-logan, cosine, hamming, hann, tukey, lanczos,
        # triangular, gaussian, barlett-hann, blackman, nuttall, blackman-harris,
        # blackman-nuttall, flat-top, kaiser, parzen

        # Create and run the algorithm object from the configuration structure
        alg_id = astra.algorithm.create(cfg)
        astra.algorithm.run(alg_id, 20)

        # Get the result
        rec = astra.data2d.get(rec_id)
        astra.algorithm.delete(alg_id)
        astra.data2d.delete(rec_id)
        astra.data2d.delete(sinogram_id)
        return rec

    def project(self, img):
        # return 1, numpy.zeros_like(img)
        proj_id = astra.create_projector('cuda', self.proj_geom, self.vol_geom)
        Spredict_id, Spredict = astra.create_sino(img, proj_id)
        astra.data2d.delete(Spredict_id)
        astra.projector.delete(proj_id)
        return Spredict

    def construct(self, sinogram):
        pass

    # import pylab
    # pylab.gray()
    # pylab.figure(1)
    # pylab.imshow(P)
    # pylab.figure(2)
    # pylab.imshow(sinogram)
    #
    # # Create a data object for the reconstruction
    # rec_id = astra.data2d.create('-vol', vol_geom)
    #
    # # create configuration
    # cfg = astra.astra_dict('FBP_CUDA')
    # cfg['ReconstructionDataId'] = rec_id
    # cfg['ProjectionDataId'] = sinogram_id
    # cfg['option'] = {'FilterType': 'Ram-Lak'}
    #
    # # possible values for FilterType:
    # # none, ram-lak, shepp-logan, cosine, hamming, hann, tukey, lanczos,
    # # triangular, gaussian, barlett-hann, blackman, nuttall, blackman-harris,
    # # blackman-nuttall, flat-top, kaiser, parzen
    #
    # # Create and run the algorithm object from the configuration structure
    # alg_id = astra.algorithm.create(cfg)
    # astra.algorithm.run(alg_id)
    #
    # # Get the result
    # rec = astra.data2d.get(rec_id)
    # pylab.figure(3)
    # pylab.imshow(rec)
    # pylab.show()
    #
    # # Clean up. Note that GPU memory is tied up in the algorithm object,
    # # and main RAM in the data objects.
    # astra.algorithm.delete(alg_id)
    # astra.data2d.delete(rec_id)
    # astra.data2d.delete(sinogram_id)
    # astra.projector.delete(proj_id)


## 640geo

if __name__ == '__main__':
    tmp_array = np.random.rand(320, 196)
    my_astra = AstraReconstruction()
    my_astra.do_fbp(tmp_array)
